#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
贵州茅台三种模型训练和预测脚本
1. XGBoost传统模型
2. LSTM深度学习模型  
3. Transformer Stock-GPT模型
"""

import pandas as pd
import numpy as np
import sys
import os
import warnings
warnings.filterwarnings('ignore')

def train_xgboost_model(stock_code, frequency):
    """训练XGBoost模型"""
    print(f"🚀 开始训练 {stock_code} 的 XGBoost 模型")
    print("=" * 50)
    
    try:
        # 读取数据
        data_file = f'stock_data_{stock_code}_{frequency}_with_indicators.csv'
        df = pd.read_csv(data_file)
        
        # 修复列名问题
        if 'date' in df.columns:
            df.rename(columns={'date': '日期'}, inplace=True)
        
        # 设置日期索引
        df['日期'] = pd.to_datetime(df['日期'])
        df.set_index('日期', inplace=True)
        
        # 准备特征和目标变量
        feature_columns = [col for col in df.columns if col not in ['code', 'close']]
        target_column = 'close'
        
        # 创建目标变量（预测下一天收盘价）
        df['target'] = df[target_column].shift(-1)
        df = df.dropna()
        
        # 准备训练数据
        X = df[feature_columns]
        y = df['target']
        
        # 分割数据（按时间序列）
        split_idx = int(0.8 * len(df))
        X_train, X_test = X[:split_idx], X[split_idx:]
        y_train, y_test = y[:split_idx], y[split_idx:]
        
        print(f"   训练集大小: {len(X_train)}")
        print(f"   测试集大小: {len(X_test)}")
        
        # 训练XGBoost模型
        import xgboost as xgb
        
        model = xgb.XGBRegressor(
            n_estimators=200,
            max_depth=6,
            learning_rate=0.1,
            random_state=42,
            n_jobs=-1
        )
        
        model.fit(X_train, y_train,
                  eval_set=[(X_test, y_test)],
                  early_stopping_rounds=20,
                  verbose=False)
        
        # 预测
        y_pred = model.predict(X_test)
        
        # 计算指标
        from sklearn.metrics import mean_squared_error, mean_absolute_error
        mse = mean_squared_error(y_test, y_pred)
        mae = mean_absolute_error(y_test, y_pred)
        rmse = np.sqrt(mse)
        
        print(f"✅ XGBoost模型训练完成")
        print(f"   RMSE: {rmse:.4f}")
        print(f"   MAE: {mae:.4f}")
        
        # 保存模型
        model_file = f'xgb_model_{stock_code}_{frequency}.json'
        model.save_model(model_file)
        print(f"   模型已保存: {model_file}")
        
        # 最新预测
        latest_data = X.iloc[-1:].values
        next_pred = model.predict(latest_data)[0]
        current_price = df['close'].iloc[-1]
        change_pct = (next_pred - current_price) / current_price * 100
        
        print(f"\\n📊 XGBoost预测结果:")
        print(f"   当前收盘价: {current_price:.2f}元")
        print(f"   预测收盘价: {next_pred:.2f}元")
        print(f"   预测涨跌: {change_pct:+.2f}%")
        
        return {
            'model': 'XGBoost',
            'rmse': rmse,
            'mae': mae,
            'current_price': current_price,
            'predicted_price': next_pred,
            'change_percent': change_pct
        }
        
    except Exception as e:
        print(f"❌ XGBoost训练失败: {e}")
        return None

def train_lstm_model(stock_code, frequency):
    """训练LSTM模型"""
    print(f"\\n🧠 开始训练 {stock_code} 的 LSTM 模型")
    print("=" * 50)
    
    try:
        # 强制使用CPU
        os.environ['CUDA_VISIBLE_DEVICES'] = ''
        
        # 导入必要库
        import tensorflow as tf
        from tensorflow import keras
        from sklearn.preprocessing import RobustScaler
        import joblib
        
        # 读取数据
        data_file = f'stock_data_{stock_code}_{frequency}_with_indicators.csv'
        df = pd.read_csv(data_file)
        
        # 选择特征列（去掉非数值列和目标列）
        feature_columns = [col for col in df.columns 
                          if col not in ['date', 'code', 'open', 'close']]
        target_columns = ['open', 'close']
        
        # 准备数据
        X_df = df[feature_columns].dropna()
        y_df = df[target_columns].dropna()
        
        # 确保数据长度一致
        min_len = min(len(X_df), len(y_df))
        X_df = X_df.iloc[:min_len]
        y_df = y_df.iloc[:min_len]
        
        print(f"   数据规模: {len(X_df)} 条记录")
        print(f"   特征数量: {len(feature_columns)}")
        
        # 归一化
        X_scaler = RobustScaler()
        y_scaler = RobustScaler()
        
        X_scaled = X_scaler.fit_transform(X_df)
        y_scaled = y_scaler.fit_transform(y_df)
        
        # 创建时间序列数据
        TIME_STEPS = 60
        def create_sequences(X, y, time_steps):
            Xs, ys = [], []
            for i in range(len(X) - time_steps):
                Xs.append(X[i:(i + time_steps)])
                ys.append(y[i + time_steps])
            return np.array(Xs), np.array(ys)
        
        X_seq, y_seq = create_sequences(X_scaled, y_scaled, TIME_STEPS)
        
        # 分割数据
        split_idx = int(0.8 * len(X_seq))
        X_train, X_test = X_seq[:split_idx], X_seq[split_idx:]
        y_train, y_test = y_seq[:split_idx], y_seq[split_idx:]
        
        print(f"   训练序列: {len(X_train)}")
        print(f"   测试序列: {len(X_test)}")
        
        # 构建LSTM模型
        model = keras.Sequential([
            keras.layers.LSTM(50, input_shape=(TIME_STEPS, len(feature_columns))),
            keras.layers.Dropout(0.2),
            keras.layers.Dense(2)  # 预测开盘价和收盘价
        ])
        
        model.compile(optimizer='adam', loss='mse', metrics=['mae'])
        
        # 训练模型
        history = model.fit(
            X_train, y_train,
            epochs=50,
            batch_size=32,
            validation_data=(X_test, y_test),
            verbose=0,
            callbacks=[
                keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)
            ]
        )
        
        # 预测
        y_pred = model.predict(X_test, verbose=0)
        
        # 反归一化
        y_test_orig = y_scaler.inverse_transform(y_test)
        y_pred_orig = y_scaler.inverse_transform(y_pred)
        
        # 计算指标
        rmse = np.sqrt(np.mean((y_test_orig - y_pred_orig) ** 2))
        mae = np.mean(np.abs(y_test_orig - y_pred_orig))
        
        print(f"✅ LSTM模型训练完成")
        print(f"   RMSE: {rmse:.4f}")
        print(f"   MAE: {mae:.4f}")
        
        # 保存模型和归一化器
        model_file = f'lstm_model_{stock_code}_{frequency}.h5'
        model.save(model_file)
        
        joblib.dump(X_scaler, f'X_scaler_{stock_code}_{frequency}.pkl')
        joblib.dump(y_scaler, f'y_scaler_{stock_code}_{frequency}.pkl')
        
        print(f"   模型已保存: {model_file}")
        
        # 最新预测
        last_sequence = X_scaled[-TIME_STEPS:].reshape(1, TIME_STEPS, -1)
        next_pred_scaled = model.predict(last_sequence, verbose=0)[0]
        next_pred = y_scaler.inverse_transform(next_pred_scaled.reshape(1, -1))[0]
        
        current_close = df['close'].iloc[-1]
        pred_open, pred_close = next_pred[0], next_pred[1]
        change_pct = (pred_close - current_close) / current_close * 100
        
        print(f"\\n📊 LSTM预测结果:")
        print(f"   当前收盘价: {current_close:.2f}元")
        print(f"   预测开盘价: {pred_open:.2f}元")
        print(f"   预测收盘价: {pred_close:.2f}元")
        print(f"   预测涨跌: {change_pct:+.2f}%")
        
        return {
            'model': 'LSTM',
            'rmse': rmse,
            'mae': mae,
            'current_price': current_close,
            'predicted_open': pred_open,
            'predicted_close': pred_close,
            'change_percent': change_pct
        }
        
    except Exception as e:
        print(f"❌ LSTM训练失败: {e}")
        import traceback
        traceback.print_exc()
        return None

def train_transformer_model(stock_code, frequency):
    """训练Transformer Stock-GPT模型"""
    print(f"\\n🤖 开始训练 {stock_code} 的 Stock-GPT Transformer 模型")
    print("=" * 50)
    
    try:
        # 使用现有的Transformer训练脚本
        import subprocess
        
        cmd = [
            'python3', 'run_technical_analysis.py', 
            'train_transformer', stock_code, 
            '-f', frequency, 
            '--force-cpu', 
            '--epochs', '20'  # 减少epoch数节省时间
        ]
        
        print("   启动Stock-GPT训练...")
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
        
        if result.returncode == 0:
            print("✅ Stock-GPT Transformer训练完成")
            
            # 尝试进行预测
            pred_cmd = [
                'python3', 'run_technical_analysis.py',
                'predict_transformer', stock_code,
                '-f', frequency,
                '--force-cpu'
            ]
            
            pred_result = subprocess.run(pred_cmd, capture_output=True, text=True, timeout=60)
            
            if pred_result.returncode == 0:
                # 从输出中解析预测结果
                output = pred_result.stdout
                print("📊 Stock-GPT预测完成")
                
                # 简单解析（实际应该更robust）
                try:
                    lines = output.split('\\n')
                    for line in lines:
                        if '当前收盘价:' in line:
                            current_price = float(line.split(':')[1].strip().replace('元', ''))
                        elif '预测开盘价:' in line:
                            pred_open = float(line.split(':')[1].strip().split()[0])
                        elif '预测收盘价:' in line:
                            pred_close = float(line.split(':')[1].strip().split()[0])
                    
                    change_pct = (pred_close - current_price) / current_price * 100
                    
                    return {
                        'model': 'Transformer',
                        'rmse': 0.0,  # 需要从训练日志中提取
                        'mae': 0.0,
                        'current_price': current_price,
                        'predicted_open': pred_open,
                        'predicted_close': pred_close,
                        'change_percent': change_pct
                    }
                except:
                    print("⚠️ 预测结果解析失败，返回基础信息")
                    return {'model': 'Transformer', 'status': 'trained'}
            else:
                print("⚠️ 预测步骤失败")
                return {'model': 'Transformer', 'status': 'trained_no_prediction'}
        else:
            print(f"❌ Stock-GPT训练失败")
            print(result.stderr)
            return None
            
    except Exception as e:
        print(f"❌ Transformer训练失败: {e}")
        return None

def main():
    """主函数：训练三种模型并对比结果"""
    stock_code = '600519'  # 贵州茅台
    frequency = 'daily'
    
    print("🍶 贵州茅台三种AI模型训练与预测")
    print("=" * 60)
    
    results = []
    
    # 1. 训练XGBoost
    xgb_result = train_xgboost_model(stock_code, frequency)
    if xgb_result:
        results.append(xgb_result)
    
    # 2. 训练LSTM
    lstm_result = train_lstm_model(stock_code, frequency)
    if lstm_result:
        results.append(lstm_result)
    
    # 3. 训练Transformer
    transformer_result = train_transformer_model(stock_code, frequency)
    if transformer_result:
        results.append(transformer_result)
    
    # 对比结果
    print("\\n🏆 三种模型预测结果对比")
    print("=" * 60)
    
    for result in results:
        if 'predicted_close' in result:
            print(f"📊 {result['model']} 模型:")
            print(f"   当前价格: {result['current_price']:.2f}元")
            if 'predicted_open' in result:
                print(f"   预测开盘: {result['predicted_open']:.2f}元")
            print(f"   预测收盘: {result.get('predicted_close', result.get('predicted_price', 0)):.2f}元")
            print(f"   预测涨跌: {result['change_percent']:+.2f}%")
            if 'rmse' in result and result['rmse'] > 0:
                print(f"   模型精度: RMSE={result['rmse']:.4f}")
            print()

if __name__ == "__main__":
    main()